#include "cartographer/mapping/internal/2d/scan_matching/adaptive_matcher.h"

namespace cartographer
{
    Eigen::Vector3d ComputeOffset(const transform::Rigid2d &pose_init, const transform::Rigid2d &pose_opt)
    {
        double init_rotation = pose_init.rotation().angle();
        const auto init_translation = pose_init.translation();

        double opt_rotation = pose_opt.rotation().angle();
        const auto opt_translation = pose_opt.translation();

        auto exp_rotation = opt_rotation - init_rotation;
        auto exp_trans = opt_translation - init_translation;
        return Eigen::Vector3d(exp_trans.x(), exp_trans.y(), exp_rotation);
    }

    float ComputeCandidateScore(const TSDF2D &tsdf, const DiscreteScan2D &discrete_scan, int x_index_offset, int y_index_offset)
    {
        float candidate_score = 0.f;
        float summed_weight = 0.f;
        for (const Eigen::Array2i &xy_index : discrete_scan)
        {
            const Eigen::Array2i proposed_xy_index(xy_index.x() + x_index_offset, xy_index.y() + y_index_offset);
            const std::pair<float, float> tsd_and_weight = tsdf.GetTSDAndWeight(proposed_xy_index);
            const float normalized_tsd_score =
                (tsdf.GetMaxCorrespondenceCost() - std::abs(tsd_and_weight.first)) / tsdf.GetMaxCorrespondenceCost();
            const float weight = tsd_and_weight.second;
            candidate_score += normalized_tsd_score * weight;
            summed_weight += weight;
        }
        if (summed_weight == 0.f)
            return 0.f;
        candidate_score /= summed_weight;
        CHECK_GE(candidate_score, 0.f);
        return candidate_score;
    }

    float ComputeCandidateScore(const ProbabilityGrid &probability_grid, const DiscreteScan2D &discrete_scan, int x_index_offset,
                                int y_index_offset)
    {
        float candidate_score = 0.f;
        for (const Eigen::Array2i &xy_index : discrete_scan)
        {
            const Eigen::Array2i proposed_xy_index(xy_index.x() + x_index_offset, xy_index.y() + y_index_offset);
            const float probability = probability_grid.GetProbability(proposed_xy_index);
            candidate_score += probability;
        }
        candidate_score /= static_cast<float>(discrete_scan.size());
        CHECK_GT(candidate_score, 0.f);
        return candidate_score;
    }

    CandidateScoreComputor2D::CandidateScoreComputor2D(const mapping::Grid2D &grid2d)
        : grid2d_(grid2d)
    {
        const auto mapMax = grid2d_.limits().max();
        LOG(INFO) << "mapMax " << mapMax(0) << " " << mapMax(1);
    }

    float CandidateScoreComputor2D::Compute(const sensor::PointCloud &rotated_cloud, int scan_index, const double x, const double y)
    {
        float score = 0.f;
        const auto &map_limits = grid2d_.limits();
        // const auto mapMax = map_limits.max();
        const int num_x_cells = grid2d_.limits().cell_limits().num_x_cells;
        const int num_y_cells = grid2d_.limits().cell_limits().num_y_cells;
        const Eigen::Affine2f translation(Eigen::Translation2f(x, y));
        // if (grid2d_.GetGridType() == GridType::PROBABILITY_GRID)
        const ProbabilityGrid &prob_grid = static_cast<const ProbabilityGrid &>(grid2d_);

        for (const auto &point : rotated_cloud)
        {
            const auto &transPt = translation * point.position.head<2>();
            // if (transPt(0) >= mapMax(0) || transPt(1) >= mapMax(1))
            //     continue;
            const auto index = map_limits.GetCellIndex(transPt);
            if (!map_limits.Contains(index))
                continue;
            // if (index(0) >= num_x_cells || index(1) >= num_y_cells)
            //     continue;
            const float probability = prob_grid.GetProbability(index);
            score += probability;
        }
        score /= static_cast<float>(rotated_cloud.size());
        // CHECK_GT(score, 0.f);
        return score;
    }

    double GaussProbability(const Eigen::Vector3d &delta, const Eigen::Vector3d &sigma)
    {
        const double ratio_x = std::pow(delta(0), 2) / sigma(0);
        const double ratio_y = std::pow(delta(1), 2) / sigma(1);
        const double ratio_yaw = std::pow(delta(2), 2) / sigma(2);
        const double num = std::exp(-0.5 * (ratio_x + ratio_y + ratio_yaw));
        const double denom = 2 * M_PI * sigma(0) * sigma(1);
        return num / denom;
    }

    bool CompareScorePair(std::pair<int, int> a, std::pair<int, int> b)
    {
        return a.second > b.second;
    }

    void PrintCandidate(const char *tag, const Candidate2D &cand)
    {
        LOG(INFO) << tag << " id " << cand.id << " x " << cand.x << " y " << cand.y << " yaw " << common::RadToDeg(cand.orientation)
                  << " score " << cand.score << " offx " << cand.x_index_offset << " offy " << cand.y_index_offset << " scan_idx "
                  << cand.scan_index;
    }

    float ComputeRangeMax(const sensor::PointCloud &cloud, const float resolution)
    {
        float max_range = 3.f * resolution;
        for (const auto &point : cloud)
        {
            const float range = point.position.head<2>().norm();
            max_range = std::max(range, max_range);
        }
        return max_range;
    }

    void NormalizeWeight(std::vector<double> &weights, const double weight_sum)
    {
        for (size_t i = 0; i < weights.size(); i++)
            weights[i] = weights[i] / weight_sum;
    }

    std::vector<sensor::PointCloud> GenerateRotatedScans(const sensor::PointCloud &cloud, const std::vector<double> &rotations)
    {
        std::vector<sensor::PointCloud> rotated_scans;
        rotated_scans.reserve(rotations.size());
        for (double rotate : rotations)
        {
            rotated_scans.push_back(
                sensor::TransformPointCloud(cloud, transform::Rigid3f::Rotation(Eigen::AngleAxisf(rotate, Eigen::Vector3f::UnitZ()))));
        }
        return rotated_scans;
    }

} // namespace cartographer
